Update init group for adaptive muon - #132

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Greptile Summary

This PR refactors AdaptiveMuon's state initialization by promoting it from ad-hoc inline logic inside step() into a proper _init_group override, enabling Megatron distributed checkpoint workflows that require pre-allocating optimizer state before gradients are present (skip_non_grad_params=False).

Key changes:

• _initialize_moment2 is removed; its logic is folded into the new _init_group override alongside momentum_buffer init, keeping both buffers in sync.
• step() now calls self._init_group(group) up-front (matching the parent class pattern), removing the "momentum_buffer" not in state guard that was previously inline.
• The normuon branch gains an explicit ndim != 2 check in _init_group, which provides a clearer error than the previous IndexError that would occur on p.data.shape[-2] for a 1-D tensor.
• The adamuon branch is missing the equivalent 2D guard — when _init_group is called externally with non-2D parameters it silently allocates partial state, sets len(state) != 0, and then step() raises a less informative error; see inline comment for details.

Confidence Score: 3/5

• Safe for normal training use, but the missing 2D guard in the adamuon branch can silently corrupt optimizer state when _init_group is called externally with non-2D parameters — the primary scenario this PR is designed to support.
• The refactoring is conceptually sound and the normuon path is correctly guarded. However, the adamuon path allocates momentum_buffer (and sets len(state) != 0) for non-2D params without raising, which breaks subsequent re-initialization attempts and defeats the purpose of the early-init API added for Megatron. This is a logic gap in the new public surface area of _init_group.
• emerging_optimizers/orthogonalized_optimizers/adaptive_muon.py — specifically the adamuon branch of _init_group (lines 130–132).

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["_init_group(group, skip_non_grad_params)"] --> B["For each param p"]
    B --> C{"skip_non_grad_params\nAND p.grad is None?"}
    C -->|Yes - skip| B
    C -->|No| D{"len(state) == 0?"}
    D -->|No - already init| B
    D -->|Yes| E["state momentum_buffer = zeros_like(p.data)"]
    E --> F{"moment2_method?"}
    F -->|adamuon| G["NO ndim check - allocates for any shape"]
    G --> H["state moment2_buffer = zeros_like(p.data)"]
    H --> B
    F -->|normuon| I{"p.data.ndim != 2?"}
    I -->|Yes| J["raise ValueError - clear early error"]
    I -->|No| K["compute moment2_shape reduced buffer"]
    K --> B
    F -->|other| L["raise TypeError"]

_{Last reviewed commit: "fix AI format error"}

/ok to test 33782da

Test Results

   48 files  ±0     98 suites  ±0   1m 15s ⏱️ -1s
1 010 tests ±0  1 010 ✅ ±0  0 💤 ±0  0 ❌ ±0 
2 251 runs  ±0  2 251 ✅ ±0  0 💤 ±0  0 ❌ ±0 

Results for commit e06ba85. ± Comparison against base commit 9c7dc7e.

♻️ This comment has been updated with latest results.

Codecov Report

❌ Patch coverage is 83.33333% with 3 lines in your changes missing coverage. Please review.

📢 Thoughts on this report? Let us know!

/ok to test 69627ba

/ok to test e06ba85